import os
import torch.utils.data as data
from utils.config import *
from utils.io_helper import read

# sgbm_root = 'G:/USVInland/Stereo Matching/dispVisualization/SGBM_Disp_Map' # Windows
# gt_root = 'G:/USVInland/Stereo Matching/Low_Res_640_320/Disp_Map'

# USVInland数据集SGBM视差验证
class USVInland_SGBM(data.Dataset):
    def __init__(self, sgbm_root, gt_root):
        sgbm_imgs = list()
        gt_imgs = list()

        # 拼接图片文件名，并加入文件名list
        for img in os.listdir(sgbm_root):
            sgbm_imgs.append(os.path.join(sgbm_root, img))
            gt_imgs.append(os.path.join(gt_root, img))

        self.sgbm_imgs = sgbm_imgs
        self.gt_imgs = gt_imgs

    def __getitem__(self, index):
        # 读取RGB图像
        if system == 'Windows': img_name = self.sgbm_imgs[index].split('\\')[-1]
        elif system == 'Ubuntu': img_name = self.sgbm_imgs[index].split('/')[-1]

        # 读取视差灰度图像
        sgbm_img = read(self.sgbm_imgs[index])
        gt_img = read(self.gt_imgs[index]) / 255 * 50

        # 数据保存到字典
        data = {'sgbm': sgbm_img, 'gt': gt_img, 'name': img_name}

        return data

    def __len__(self):
        return len(self.sgbm_imgs)
